#include "printf.h"
#include "pmm.h"
#include "vmm.h"
#include "interrupt.h"
#include "clint.h"
#include "riscv.h"
#include "mmio.h"
#include "sbi.h"        // SBI接口
#include "scheduler.h"   // 调度器
#include <stdint.h>
#include <stddef.h>

// --------------------------
// UART寄存器定义
// --------------------------
#define UART0_ADDR 0x10000000UL
#define UART_RHR (UART0_ADDR + 0x0)
#define UART_THR (UART0_ADDR + 0x0)
#define UART_LSR (UART0_ADDR + 0x5)
#define UART_IER (UART0_ADDR + 0x1)
#define UART_IIR (UART0_ADDR + 0x2)

#define UART_LSR_RX_READY (1 << 0)
#define UART_LSR_TX_EMPTY (1 << 5)
#define UART_IER_RX_INT (1 << 0)

// --------------------------
// 延时函数
// --------------------------
void delay(int count) {
    for (int i = 0; i < count; i++) {
        for (int j = 0; j < 1000; j++) {
            __asm__ __volatile__("nop");
        }
    }
}

// --------------------------
// 定时器中断处理函数(支持调度和测试模式)
// --------------------------
static volatile int timer_irq_count = 0;
static volatile int test_mode = 0;

void timer_handler(void) {
    timer_irq_count++;
    
    // 设置下次中断时间
    uint64_t next_time = get_time() + 1000000;
    sbi_set_timer(next_time);
    
    // 测试模式
    if (test_mode) {
        if (timer_irq_count <= 5) {
            printf("[中断测试] 第%d次时钟中断\n", timer_irq_count);
        }
        
        if (timer_irq_count == 5) {
            test_mode = 0;
            printf("[中断测试] 测试完成!\n\n");
        }
        return;
    }
    
    // 正常模式: 调度逻辑
    task_t* current = get_current_task();
    
    if (current) {
        // 更新任务时间统计
        current->elapsed_time += 1;
        
        // 定期打印时钟中断信息
        if (timer_irq_count % 1000 == 0) {
            printf("\n[时钟中断] 第%d次触发(约%d秒), 当前任务: %s(PID=%d)\n",
                   timer_irq_count, timer_irq_count / 1000,
                   current->name, current->pid);
            print_tasks();
        }
        
        // 检查时间片是否用完
        if (current->elapsed_time >= current->time_slice) {
            printf("[调度] 任务%s时间片耗尽(%dms/%dms),触发调度\n",
                   current->name, current->elapsed_time, current->time_slice);
            schedule();
        }
    }
}

// --------------------------
// 中断功能测试函数
// --------------------------
void test_timer_interrupt(void) {
    printf("\n=== 中断功能测试 ===\n");
    printf("Testing timer interrupt...\n");
    
    timer_irq_count = 0;
    test_mode = 1;
    
    printf("等待5次中断...\n");
    
    while (test_mode) {
        wfi();
    }
    
    printf("Timer test completed: 5 interrupts\n");
    printf("=== 中断功能测试完成 ===\n\n");
}

// --------------------------
// UART功能函数
// --------------------------
static int uart_getc(void) {
    if (mmio_read32(UART_LSR) & UART_LSR_RX_READY) {
        return mmio_read32(UART_RHR) & 0xFF;
    } else {
        return -1;
    }
}

static void uart_init_irq(void) {
    mmio_write32(UART_IER, UART_IER_RX_INT);
}

static void uart_clear_irq(void) {
    mmio_read32(UART_IIR);
}

// --------------------------
// UART中断处理函数
// --------------------------
static volatile int uart_irq_count = 0;

void uart_handler(void) {
    uart_irq_count++;
    int c = uart_getc();
    if (c != -1) {
        printf("\n[UART中断] 第%d次触发，收到字符: '%c' (ASCII: %d)\n",
               uart_irq_count, c, c);
        uart_putc(c);
    }
    uart_clear_irq();
}

// --------------------------
// 虚拟地址访问测试
// --------------------------
void test_virtual_address_access() {
    void* phys_page = alloc_page();
    if (!phys_page) {
        printf("test_virtual_address_access: 分配物理页失败\n");
        return;
    }
    printf("\n=== 虚拟地址访问测试 ===\n");
    uint64_t pa = (uint64_t)phys_page;
    uint64_t va = KERNBASE + (pa - PHYSBASE);
    printf("物理页: PA = %p -> 虚拟地址: VA = %p\n", phys_page, (void*)va);
    printf("=== 虚拟地址访问测试结束 ===\n");
}

// --------------------------
// 测试任务函数
// --------------------------
void task_a(void* arg) {
    (void)arg;
    int count = 0;
    while (1) {
        printf("[任务A] 正在执行...计数=%d\n", count++);
        delay(500);
        if (count >= 10) count = 0;
    }
}

void task_b(void* arg) {
    (void)arg;
    int count = 0;
    while (1) {
        printf("[任务B] 正在执行...计数=%d\n", count++);
        delay(300);
        if (count >= 10) count = 0;
    }
}

void task_c(void* arg) {
    (void)arg;
    int count = 0;
    while (1) {
        printf("[任务C] 正在执行...计数=%d\n", count++);
        delay(400);
        if (count >= 10) count = 0;
    }
}

// --------------------------
// 主函数
// --------------------------
int main(void) {
    // 1. 启动信息与printf测试
    printf("RISC-V 操作系统启动...\n");
    printf("Hello, RISC-V OS!\n");
    printf("Decimal test: %d, %d, %d\n", 123, -456, 0);
    printf("Hex test: %x, %x\n", 0x1234, 0xABCD);
    printf("String test: %s\n", "This is a string");
    printf("Char test: %c, %c\n", 'A', 'Z');
    printf("Mixed test: %s = %d (0x%x)\n", "Answer", 42, 42);
    printf("Percent test: 100%% complete\n");
    
    // 2. 清屏测试
    printf("3秒后清屏...\n");
    delay(3000);
    clear_screen();
    printf("屏幕已清空！\n");
    
    // 3. 物理内存管理测试
    printf("\n=== 物理内存管理测试 ===\n");
    pmm_init((uint64_t)_end, PHYSTOP);
    printf("PMM初始化完成（内存范围：%p ~ %p）\n", (void*)_end, (void*)PHYSTOP);
    
    void* page1 = alloc_page();
    printf("分配单个页: %p %s\n", page1, page1 ? "（成功）" : "（失败）");
    
    void* pages = alloc_pages(5);
    printf("分配5个连续页: %p %s\n", pages, pages ? "（成功）" : "（失败）");
    
    size_t total, free;
    get_memory_info(&total, &free);
    printf("内存状态: 总内存 %d KB, 空闲内存 %d KB\n", total/1024, free/1024);
    
    if (page1) free_page(page1);
    printf("释放单个页: %p（完成）\n", page1);
    if (pages) free_pages(pages, 5);
    printf("释放5个连续页: %p（完成）\n", pages);
    
    get_memory_info(&total, &free);
    printf("释放后内存状态: 总内存 %d KB, 空闲内存 %d KB\n", total/1024, free/1024);
    
    void* large_alloc = alloc_pages(100);
    if (large_alloc) {
        printf("分配100个连续页: %p（成功）\n", large_alloc);
        free_pages(large_alloc, 100);
    } else {
        printf("分配100个连续页: 失败\n");
    }
    printf("=== 物理内存管理测试结束 ===\n");
    
    // 4. 虚拟内存初始化
    printf("\n=== 虚拟内存初始化 ===\n");
    kvm_init();
    printf("=== 虚拟内存已启用 ===\n");
    
    // 5. 内核页表映射验证
    printf("\n=== 内核页表映射验证 ===\n");
    pte_t* text_pte = walk_lookup(kernel_pagetable, (uint64_t)_text + (KERNBASE - PHYSBASE));
    if (text_pte && (*text_pte & (PTE_V | PTE_R | PTE_X)) == (PTE_V | PTE_R | PTE_X)) {
        printf("代码段映射验证: 成功（R+X权限）\n");
    } else {
        printf("代码段映射验证: 失败\n");
    }
    
    pte_t* uart_pte = walk_lookup(kernel_pagetable, UART0);
    if (uart_pte && (*uart_pte & (PTE_V | PTE_R | PTE_W)) == (PTE_V | PTE_R | PTE_W)) {
        printf("UART设备映射验证: 成功（R+W权限）\n");
    } else {
        printf("UART设备映射验证: 失败\n");
    }
    printf("=== 内核页表映射验证结束 ===\n");
    
    // 5.2 虚拟地址读写测试
    test_virtual_address_access();
    
    // 5.3 打印页表结构
    printf("\n=== 内核页表结构 ===\n");
    dump_pagetable(kernel_pagetable, 0);
    
    // 6. 中断处理测试
    printf("\n=== 中断处理测试 ===\n");
    trap_init();
    printf("中断系统初始化完成\n");
    
    // 6.1 注册并启用时钟中断
    int timer_reg = register_interrupt(IRQ_TIMER, timer_handler, IRQ_PRIORITY_3);
    if (timer_reg != 0) {
        printf("注册时钟中断失败！错误码: %d\n", timer_reg);
    } else {
        enable_interrupt(IRQ_TIMER);
        
        uint64_t first_timer = get_time() + 1000000;
        sbi_set_timer(first_timer);
        
        printf("时钟中断已启用（每1ms触发一次）\n");
    }
    
    // 6.1.5 执行中断功能测试
    printf("\n开始执行中断功能测试...\n");
    test_timer_interrupt();
    printf("中断测试已完成，继续正常运行...\n");
    
    // 重置中断计数器
    timer_irq_count = 0;
    
    // 6.2 注册并启用UART中断
    int uart_reg = register_interrupt(IRQ_UART, uart_handler, IRQ_PRIORITY_5);
    if (uart_reg != 0) {
        printf("注册UART中断失败！错误码: %d\n", uart_reg);
    } else {
        uart_init_irq();
        enable_interrupt(IRQ_UART);
        printf("UART中断已启用（输入字符可触发）\n");
    }
    printf("=== 中断处理测试启动完成 ===\n");
    
    // 7. 调度器初始化(现在启用)
    printf("\n=== 调度器初始化 ===\n");
    scheduler_init();
    
    // 创建测试任务
    create_task("TaskA", task_a, NULL, 50);
    create_task("TaskB", task_b, NULL, 30);
    create_task("TaskC", task_c, NULL, 40);
    
    printf("\n=== 初始任务列表 ===\n");
    print_tasks();
    printf("=== 调度器启动完成 ===\n");
    
    // 8. 内核主循环
    printf("\n内核运行中...（支持多任务调度）\n");
    printf("提示: 输入字符测试UART中断，按Ctrl+C退出QEMU\n\n");
    
    while (1) {
        wfi();
        
        // 周期性打印运行状态
        static int loop_count = 0;
        if (++loop_count % 5000 == 0) {
            task_t* current = get_current_task();
            printf("[主循环] 循环计数: %d, 当前任务: %s, 时钟中断: %d次, UART中断: %d次\n",
                   loop_count, current ? current->name : "无",
                   timer_irq_count, uart_irq_count);
        }
    }
    
    return 0;
}

